{
    "jax": {
        "classes": {
            "Array": {
                "Array": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ShapeDtypeStruct": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__len__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Shard": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Device": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShapeDtypeStruct": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Shard": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_ShapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "custom_jvp": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "custom_vjp": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "__getattr__",
            "_deprecated_curry",
            "block_until_ready",
            "checkpoint",
            "clear_backends",
            "closure_convert",
            "custom_gradient",
            "default_backend",
            "device_count",
            "device_get",
            "device_put",
            "device_put_replicated",
            "device_put_sharded",
            "devices",
            "disable_jit",
            "effects_barrier",
            "ensure_compile_time_eval",
            "eval_shape",
            "grad",
            "hessian",
            "host_count",
            "host_id",
            "host_ids",
            "jacfwd",
            "jacobian",
            "jacrev",
            "jit",
            "jvp",
            "linear_transpose",
            "linearize",
            "live_arrays",
            "local_device_count",
            "local_devices",
            "make_array_from_callback",
            "make_array_from_single_device_arrays",
            "make_jaxpr",
            "named_call",
            "named_scope",
            "pmap",
            "print_environment_info",
            "process_count",
            "process_index",
            "pure_callback",
            "remat",
            "transfer_guard",
            "tree_flatten",
            "tree_leaves",
            "tree_map",
            "tree_structure",
            "tree_transpose",
            "tree_unflatten",
            "treedef_is_leaf",
            "value_and_grad",
            "vjp",
            "vmap",
            "xla_computation"
        ],
        "methods": [],
        "modules": {
            "jax.lax": {
                "classes": {
                    "ConvDimensionNumbers": {
                        "ConvDimensionNumbers": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "GatherDimensionNumbers": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "GatherScatterMode": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [],
                            "methods": [],
                            "modules": {}
                        },
                        "Precision": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "RandomAlgorithm": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "RoundingMethod": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [],
                            "methods": [],
                            "modules": {}
                        },
                        "ScatterDimensionNumbers": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "GatherDimensionNumbers": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "GatherScatterMode": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Precision": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "RandomAlgorithm": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "RoundingMethod": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ScatterDimensionNumbers": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "abs",
                    "acos",
                    "acosh",
                    "add",
                    "after_all",
                    "all_gather",
                    "all_to_all",
                    "approx_max_k",
                    "approx_min_k",
                    "argmax",
                    "argmin",
                    "asin",
                    "asinh",
                    "associative_scan",
                    "atan",
                    "atan2",
                    "atanh",
                    "axis_index",
                    "batch_matmul",
                    "bessel_i0e",
                    "bessel_i1e",
                    "betainc",
                    "bitcast_convert_type",
                    "bitwise_and",
                    "bitwise_not",
                    "bitwise_or",
                    "bitwise_xor",
                    "broadcast",
                    "broadcast_in_dim",
                    "broadcast_shapes",
                    "broadcast_to_rank",
                    "broadcasted_iota",
                    "cbrt",
                    "ceil",
                    "clamp",
                    "clz",
                    "collapse",
                    "complex",
                    "concatenate",
                    "cond",
                    "conj",
                    "conv",
                    "conv_dimension_numbers",
                    "conv_general_dilated",
                    "conv_general_dilated_local",
                    "conv_general_dilated_patches",
                    "conv_general_permutations",
                    "conv_general_shape_tuple",
                    "conv_shape_tuple",
                    "conv_transpose",
                    "conv_transpose_shape_tuple",
                    "conv_with_general_padding",
                    "convert_element_type",
                    "cos",
                    "cosh",
                    "create_token",
                    "cumlogsumexp",
                    "cummax",
                    "cummin",
                    "cumprod",
                    "cumsum",
                    "custom_linear_solve",
                    "custom_root",
                    "digamma",
                    "div",
                    "dot",
                    "dot_general",
                    "dynamic_index_in_dim",
                    "dynamic_slice",
                    "dynamic_slice_in_dim",
                    "dynamic_update_index_in_dim",
                    "dynamic_update_slice",
                    "dynamic_update_slice_in_dim",
                    "eq",
                    "erf",
                    "erf_inv",
                    "erfc",
                    "exp",
                    "expand_dims",
                    "expm1",
                    "floor",
                    "fori_loop",
                    "full",
                    "full_like",
                    "gather",
                    "ge",
                    "gt",
                    "igamma",
                    "igamma_grad_a",
                    "igammac",
                    "imag",
                    "index_in_dim",
                    "index_take",
                    "infeed",
                    "integer_pow",
                    "iota",
                    "is_finite",
                    "le",
                    "lgamma",
                    "log",
                    "log1p",
                    "logistic",
                    "lt",
                    "map",
                    "max",
                    "min",
                    "mul",
                    "naryop",
                    "naryop_dtype_rule",
                    "ne",
                    "neg",
                    "nextafter",
                    "outfeed",
                    "pad",
                    "padtype_to_pads",
                    "pdot",
                    "pmax",
                    "pmean",
                    "pmin",
                    "population_count",
                    "pow",
                    "ppermute",
                    "prod",
                    "pshuffle",
                    "psum",
                    "psum_scatter",
                    "pswapaxes",
                    "random_gamma_grad",
                    "real",
                    "reciprocal",
                    "reduce",
                    "reduce_precision",
                    "reduce_window",
                    "reduce_window_shape_tuple",
                    "rem",
                    "reshape",
                    "rev",
                    "rng_bit_generator",
                    "rng_uniform",
                    "round",
                    "rsqrt",
                    "scan",
                    "scan_bind",
                    "scatter",
                    "scatter_add",
                    "scatter_apply",
                    "scatter_max",
                    "scatter_min",
                    "scatter_mul",
                    "select",
                    "select_n",
                    "shift_left",
                    "shift_right_arithmetic",
                    "shift_right_logical",
                    "sign",
                    "sin",
                    "sinh",
                    "slice",
                    "slice_in_dim",
                    "sort",
                    "sort_key_val",
                    "sqrt",
                    "square",
                    "squeeze",
                    "standard_abstract_eval",
                    "standard_primitive",
                    "stop_gradient",
                    "sub",
                    "switch",
                    "tan",
                    "tanh",
                    "tie_in",
                    "top_k",
                    "transpose",
                    "unop",
                    "unop_dtype_rule",
                    "while_loop",
                    "with_sharding_constraint",
                    "xeinsum",
                    "zeros_like_array"
                ],
                "methods": [],
                "modules": {
                    "jax.lax.linalg": {
                        "classes": {},
                        "functions": [
                            "cholesky",
                            "eig",
                            "eigh",
                            "hessenberg",
                            "householder_product",
                            "lu",
                            "lu_pivots_to_permutation",
                            "qr",
                            "schur",
                            "svd",
                            "triangular_solve",
                            "tridiagonal",
                            "tridiagonal_solve"
                        ],
                        "methods": [],
                        "modules": {}
                    }
                }
            },
            "jax.lib": {
                "classes": {},
                "functions": [],
                "methods": [],
                "modules": {
                    "jax.lib.xla_bridge": {
                        "classes": {},
                        "functions": [
                            "default_backend",
                            "get_compile_options"
                        ],
                        "methods": [],
                        "modules": {}
                    }
                }
            },
            "jax.linear_util": {
                "classes": {
                    "StoreException": {
                        "StoreException": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "WrappedFun": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "call_wrapped",
                                "populate_stores",
                                "wrap"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "WrappedFun": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "cache",
                    "merge_linear_aux",
                    "wrap_init"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.monitoring": {
                "classes": {},
                "functions": [
                    "record_event",
                    "record_event_duration_secs",
                    "register_event_duration_secs_listener",
                    "register_event_listener"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.nn": {
                "classes": {},
                "functions": [
                    "gelu",
                    "logsumexp",
                    "normalize",
                    "one_hot",
                    "softmax"
                ],
                "methods": [],
                "modules": {
                    "jax.nn.initializers": {
                        "classes": {
                            "Initializer": {
                                "Initializer": {
                                    "classes": {
                                        "__class__": {
                                            "classes": {},
                                            "functions": [],
                                            "methods": []
                                        }
                                    },
                                    "functions": [
                                        "__call__",
                                        "__init__",
                                        "__new__",
                                        "__subclasshook__"
                                    ],
                                    "methods": [
                                        "__class_getitem__",
                                        "__init_subclass__"
                                    ],
                                    "modules": {}
                                },
                                "classes": {},
                                "functions": [],
                                "methods": []
                            }
                        },
                        "functions": [
                            "constant",
                            "delta_orthogonal",
                            "glorot_normal",
                            "glorot_uniform",
                            "he_normal",
                            "he_uniform",
                            "kaiming_normal",
                            "kaiming_uniform",
                            "lecun_normal",
                            "lecun_uniform",
                            "normal",
                            "ones",
                            "orthogonal",
                            "uniform",
                            "variance_scaling",
                            "xavier_normal",
                            "xavier_uniform",
                            "zeros"
                        ],
                        "methods": [],
                        "modules": {}
                    }
                }
            },
            "jax.numpy": {
                "classes": {
                    "ComplexWarning": {
                        "ComplexWarning": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "bool_": {
                            "classes": {
                                "__class__": {
                                    "_ScalarMeta": {
                                        "classes": {
                                            "__base__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            },
                                            "__class__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            }
                                        },
                                        "functions": [
                                            "__call__",
                                            "__eq__",
                                            "__hash__",
                                            "__init_subclass__",
                                            "__instancecheck__",
                                            "__ne__",
                                            "__new__",
                                            "__prepare__",
                                            "__subclasshook__"
                                        ],
                                        "methods": [],
                                        "modules": {}
                                    },
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "character": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "complex128": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "complex64": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "complexfloating": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "dtype": {
                            "classes": {
                                "__class__": {
                                    "_DTypeMeta": {
                                        "classes": {
                                            "__base__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            },
                                            "__class__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            }
                                        },
                                        "functions": [
                                            "__init_subclass__",
                                            "__new__",
                                            "__prepare__",
                                            "__subclasshook__"
                                        ],
                                        "methods": [],
                                        "modules": {}
                                    },
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "flexible": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "float16": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "float32": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "float64": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "float8_e4m3fn": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "float8_e5m2": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "floating": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "functions": [],
                        "generic": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "inexact": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "int16": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "int32": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "int64": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "int8": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "integer": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "methods": [],
                        "number": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "object_": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "signedinteger": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "uint16": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "uint32": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "uint64": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "uint8": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "unsignedinteger": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        }
                    },
                    "bfloat16": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "bool_": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "cdouble": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "character": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "complex128": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "complex64": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "complex_": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "complexfloating": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "csingle": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "double": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "dtype": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "finfo": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "flexible": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float16": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float32": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float64": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float8_e4m3fn": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float8_e5m2": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "float_": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "floating": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "generic": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "iinfo": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "inexact": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "int16": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "int32": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "int64": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "int8": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "int_": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "integer": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ndarray": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "number": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "object_": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "signedinteger": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "single": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "uint": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "uint16": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "uint32": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "uint64": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "uint8": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "unsignedinteger": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "__getattr__",
                    "all",
                    "alltrue",
                    "amax",
                    "amin",
                    "any",
                    "apply_along_axis",
                    "apply_over_axes",
                    "arange",
                    "argmax",
                    "argmin",
                    "argwhere",
                    "array",
                    "array_equal",
                    "array_equiv",
                    "array_repr",
                    "array_split",
                    "array_str",
                    "asarray",
                    "average",
                    "bartlett",
                    "bincount",
                    "blackman",
                    "broadcast_arrays",
                    "broadcast_shapes",
                    "broadcast_to",
                    "can_cast",
                    "choose",
                    "column_stack",
                    "compress",
                    "concatenate",
                    "copy",
                    "cumprod",
                    "cumproduct",
                    "cumsum",
                    "delete",
                    "diag",
                    "diag_indices",
                    "diag_indices_from",
                    "diagflat",
                    "dsplit",
                    "dstack",
                    "einsum",
                    "einsum_path",
                    "empty",
                    "empty_like",
                    "expand_dims",
                    "extract",
                    "eye",
                    "flatnonzero",
                    "flip",
                    "fliplr",
                    "flipud",
                    "from_dlpack",
                    "frombuffer",
                    "fromfile",
                    "fromfunction",
                    "fromiter",
                    "fromstring",
                    "full",
                    "full_like",
                    "geomspace",
                    "get_printoptions",
                    "hamming",
                    "hanning",
                    "histogram",
                    "histogram2d",
                    "histogram_bin_edges",
                    "histogramdd",
                    "hsplit",
                    "hstack",
                    "identity",
                    "in1d",
                    "indices",
                    "insert",
                    "intersect1d",
                    "iscomplexobj",
                    "isin",
                    "isneginf",
                    "isposinf",
                    "isrealobj",
                    "isscalar",
                    "issubdtype",
                    "issubsctype",
                    "iterable",
                    "ix_",
                    "kaiser",
                    "linspace",
                    "load",
                    "logspace",
                    "mask_indices",
                    "max",
                    "mean",
                    "meshgrid",
                    "min",
                    "moveaxis",
                    "msort",
                    "nanargmax",
                    "nanargmin",
                    "nancumprod",
                    "nancumsum",
                    "ndim",
                    "nonzero",
                    "ones",
                    "ones_like",
                    "pad",
                    "piecewise",
                    "place",
                    "polydiv",
                    "polymul",
                    "power",
                    "printoptions",
                    "prod",
                    "product",
                    "promote_types",
                    "ptp",
                    "put",
                    "ravel_multi_index",
                    "repeat",
                    "reshape",
                    "result_type",
                    "roll",
                    "roots",
                    "row_stack",
                    "save",
                    "savez",
                    "select",
                    "set_printoptions",
                    "setdiff1d",
                    "setxor1d",
                    "shape",
                    "size",
                    "sometrue",
                    "split",
                    "squeeze",
                    "stack",
                    "std",
                    "sum",
                    "take",
                    "tensordot",
                    "tile",
                    "transpose",
                    "tri",
                    "tril_indices",
                    "tril_indices_from",
                    "trim_zeros",
                    "triu_indices",
                    "triu_indices_from",
                    "union1d",
                    "unique",
                    "unravel_index",
                    "var",
                    "vectorize",
                    "vsplit",
                    "vstack",
                    "where",
                    "zeros",
                    "zeros_like"
                ],
                "methods": [],
                "modules": {
                    "jax.numpy.fft": {
                        "classes": {},
                        "functions": [
                            "fft",
                            "fft2",
                            "fftfreq",
                            "fftn",
                            "fftshift",
                            "hfft",
                            "ifft",
                            "ifft2",
                            "ifftn",
                            "ifftshift",
                            "ihfft",
                            "irfft",
                            "irfft2",
                            "irfftn",
                            "rfft",
                            "rfft2",
                            "rfftfreq",
                            "rfftn"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.numpy.linalg": {
                        "classes": {},
                        "functions": [
                            "cond",
                            "eig",
                            "lstsq",
                            "multi_dot",
                            "tensorinv",
                            "tensorsolve"
                        ],
                        "methods": [],
                        "modules": {}
                    }
                }
            },
            "jax.ops": {
                "classes": {},
                "functions": [
                    "segment_max",
                    "segment_min",
                    "segment_prod",
                    "segment_sum"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.profiler": {
                "classes": {
                    "StepTraceAnnotation": {
                        "StepTraceAnnotation": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "is_enabled"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "TraceAnnotation": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "is_enabled"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "TraceAnnotation": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "annotate_function",
                    "device_memory_profile",
                    "save_device_memory_profile",
                    "start_server",
                    "start_trace",
                    "stop_server",
                    "stop_trace",
                    "trace"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.random": {
                "classes": {
                    "KeyArray": {
                        "PRNGKeyArray": {
                            "classes": {
                                "__class__": {
                                    "PRNGKeyArrayMeta": {
                                        "classes": {
                                            "__base__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            },
                                            "__class__": {
                                                "classes": {},
                                                "functions": [],
                                                "methods": []
                                            }
                                        },
                                        "functions": [
                                            "__init_subclass__",
                                            "__instancecheck__",
                                            "__new__",
                                            "__prepare__",
                                            "__subclasscheck__",
                                            "__subclasshook__",
                                            "_abc_caches_clear",
                                            "_abc_registry_clear",
                                            "_dump_registry",
                                            "register"
                                        ],
                                        "methods": [],
                                        "modules": {}
                                    },
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__getitem__",
                                "__init__",
                                "__init_subclass__",
                                "__iter__",
                                "__len__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "_chunk_iter",
                                "_is_scalar",
                                "block_until_ready",
                                "broadcast_to",
                                "concatenate",
                                "expand_dims",
                                "pprint",
                                "reshape",
                                "squeeze",
                                "take",
                                "transpose",
                                "unsafe_raw_array"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PRNGKeyArray": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "_PRNGKeyArray": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "PRNGKey",
                    "ball",
                    "bernoulli",
                    "beta",
                    "categorical",
                    "cauchy",
                    "chisquare",
                    "choice",
                    "default_prng_impl",
                    "dirichlet",
                    "double_sided_maxwell",
                    "exponential",
                    "f",
                    "fold_in",
                    "gamma",
                    "generalized_normal",
                    "gumbel",
                    "key_data",
                    "laplace",
                    "loggamma",
                    "logistic",
                    "maxwell",
                    "multivariate_normal",
                    "normal",
                    "orthogonal",
                    "pareto",
                    "permutation",
                    "poisson",
                    "rademacher",
                    "randint",
                    "rayleigh",
                    "rbg_key",
                    "shuffle",
                    "split",
                    "t",
                    "threefry2x32_key",
                    "threefry_2x32",
                    "truncated_normal",
                    "uniform",
                    "unsafe_rbg_key",
                    "wald",
                    "weibull_min"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.scipy": {
                "classes": {},
                "functions": [],
                "methods": [],
                "modules": {
                    "jax.scipy.cluster": {
                        "classes": {},
                        "functions": [],
                        "methods": [],
                        "modules": {
                            "jax.scipy.cluster.vq": {
                                "classes": {},
                                "functions": [
                                    "vq"
                                ],
                                "methods": [],
                                "modules": {}
                            }
                        }
                    },
                    "jax.scipy.fft": {
                        "classes": {},
                        "functions": [
                            "dct",
                            "dctn"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.interpolate": {
                        "classes": {
                            "RegularGridInterpolator": {
                                "RegularGridInterpolator": {
                                    "classes": {
                                        "__class__": {
                                            "classes": {},
                                            "functions": [],
                                            "methods": []
                                        },
                                        "__np_wrapped__": {
                                            "classes": {},
                                            "functions": [],
                                            "methods": []
                                        }
                                    },
                                    "functions": [
                                        "__call__",
                                        "__init__",
                                        "__init_subclass__",
                                        "__new__",
                                        "__subclasshook__",
                                        "_evaluate_linear",
                                        "_evaluate_nearest",
                                        "_find_indices"
                                    ],
                                    "methods": [],
                                    "modules": {}
                                },
                                "classes": {},
                                "functions": [],
                                "methods": []
                            }
                        },
                        "functions": [],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.linalg": {
                        "classes": {},
                        "functions": [
                            "cho_factor",
                            "cho_solve",
                            "cholesky",
                            "det",
                            "eigh",
                            "funm",
                            "inv",
                            "qr",
                            "schur",
                            "solve",
                            "solve_triangular",
                            "sqrtm",
                            "svd",
                            "toeplitz",
                            "tril",
                            "triu"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.ndimage": {
                        "classes": {},
                        "functions": [
                            "map_coordinates"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.signal": {
                        "classes": {},
                        "functions": [
                            "convolve",
                            "convolve2d",
                            "correlate",
                            "correlate2d",
                            "csd",
                            "detrend",
                            "istft",
                            "stft",
                            "welch"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.sparse": {
                        "classes": {},
                        "functions": [],
                        "methods": [],
                        "modules": {
                            "jax.scipy.sparse.linalg": {
                                "classes": {},
                                "functions": [
                                    "bicgstab",
                                    "cg",
                                    "gmres"
                                ],
                                "methods": [],
                                "modules": {}
                            }
                        }
                    },
                    "jax.scipy.special": {
                        "classes": {},
                        "functions": [
                            "betainc",
                            "betaln",
                            "digamma",
                            "entr",
                            "erf",
                            "erfc",
                            "erfinv",
                            "exp1",
                            "expit",
                            "gammainc",
                            "gammaincc",
                            "gammaln",
                            "i0",
                            "i0e",
                            "i1",
                            "i1e",
                            "logsumexp",
                            "lpmn",
                            "lpmn_values",
                            "multigammaln",
                            "ndtr",
                            "ndtri",
                            "polygamma",
                            "spence",
                            "sph_harm",
                            "xlog1py",
                            "xlogy",
                            "zeta"
                        ],
                        "methods": [],
                        "modules": {}
                    },
                    "jax.scipy.stats": {
                        "classes": {
                            "gaussian_kde": {
                                "classes": {},
                                "functions": [],
                                "gaussian_kde": {
                                    "classes": {
                                        "__class__": {
                                            "classes": {},
                                            "functions": [],
                                            "methods": []
                                        },
                                        "__np_wrapped__": {
                                            "classes": {},
                                            "functions": [],
                                            "methods": []
                                        }
                                    },
                                    "functions": [
                                        "__call__",
                                        "__delattr__",
                                        "__eq__",
                                        "__hash__",
                                        "__init__",
                                        "__init_subclass__",
                                        "__new__",
                                        "__repr__",
                                        "__setattr__",
                                        "__subclasshook__",
                                        "_reshape_points",
                                        "_setattr",
                                        "evaluate",
                                        "integrate_box",
                                        "integrate_box_1d",
                                        "integrate_gaussian",
                                        "integrate_kde",
                                        "logpdf",
                                        "pdf",
                                        "resample",
                                        "set_bandwidth",
                                        "tree_flatten"
                                    ],
                                    "methods": [
                                        "tree_unflatten"
                                    ],
                                    "modules": {}
                                },
                                "methods": []
                            }
                        },
                        "functions": [],
                        "methods": [],
                        "modules": {
                            "jax.scipy.stats.bernoulli": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logpmf",
                                    "pmf",
                                    "ppf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.beta": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logcdf",
                                    "logpdf",
                                    "pdf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.betabinom": {
                                "classes": {},
                                "functions": [
                                    "logpmf",
                                    "pmf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.cauchy": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "isf",
                                    "logcdf",
                                    "logpdf",
                                    "pdf",
                                    "ppf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.chi2": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logcdf",
                                    "logpdf",
                                    "pdf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.dirichlet": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.expon": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.gamma": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logcdf",
                                    "logpdf",
                                    "pdf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.gennorm": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.geom": {
                                "classes": {},
                                "functions": [
                                    "logpmf",
                                    "pmf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.laplace": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.logistic": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "isf",
                                    "logpdf",
                                    "pdf",
                                    "ppf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.multinomial": {
                                "classes": {},
                                "functions": [
                                    "logpmf",
                                    "pmf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.multivariate_normal": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.nbinom": {
                                "classes": {},
                                "functions": [
                                    "logpmf",
                                    "pmf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.norm": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "isf",
                                    "logcdf",
                                    "logpdf",
                                    "pdf",
                                    "ppf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.pareto": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.poisson": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logpmf",
                                    "pmf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.t": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.truncnorm": {
                                "classes": {},
                                "functions": [
                                    "cdf",
                                    "logcdf",
                                    "logpdf",
                                    "logsf",
                                    "pdf",
                                    "sf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.uniform": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "jax.scipy.stats.vonmises": {
                                "classes": {},
                                "functions": [
                                    "logpdf",
                                    "pdf"
                                ],
                                "methods": [],
                                "modules": {}
                            }
                        }
                    }
                }
            },
            "jax.sharding": {
                "classes": {
                    "GSPMDSharding": {
                        "GSPMDSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init_subclass__",
                                "__new__",
                                "__reduce__",
                                "__repr__",
                                "__subclasshook__",
                                "_to_xla_op_sharding",
                                "is_compatible_aval",
                                "is_equivalent_to"
                            ],
                            "methods": [
                                "get_replicated"
                            ],
                            "modules": {}
                        },
                        "NamedSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init_subclass__",
                                "__new__",
                                "__reduce__",
                                "__repr__",
                                "__subclasshook__",
                                "_preprocess",
                                "is_compatible_aval",
                                "is_equivalent_to"
                            ],
                            "methods": [
                                "_from_parsed_pspec"
                            ],
                            "modules": {}
                        },
                        "PmapSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init_subclass__",
                                "__new__",
                                "__reduce__",
                                "__repr__",
                                "__str__",
                                "__subclasshook__",
                                "_to_xla_op_sharding",
                                "is_equivalent_to"
                            ],
                            "methods": [
                                "default"
                            ],
                            "modules": {}
                        },
                        "PositionalSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "is_equivalent_to",
                                "replicate",
                                "reshape",
                                "transpose"
                            ],
                            "methods": [
                                "remake"
                            ],
                            "modules": {}
                        },
                        "Sharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "devices_indices_map",
                                "is_equivalent_to",
                                "shard_shape"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "SingleDeviceSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__hash__",
                                "__init_subclass__",
                                "__new__",
                                "__reduce__",
                                "__repr__",
                                "__subclasshook__",
                                "_to_xla_op_sharding",
                                "devices_indices_map",
                                "is_equivalent_to"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "XLACompatibleSharding": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "_to_xla_op_sharding",
                                "is_equivalent_to"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Mesh": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "NamedSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PartitionSpec": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PmapSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PositionalSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Sharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "SingleDeviceSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "XLACompatibleSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "_deprecated_OpShardingSharding": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "__getattr__"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.stages": {
                "classes": {
                    "ArgInfo": {
                        "ArgInfo": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "Compiled": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__call__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "as_text",
                                "call",
                                "compiler_ir",
                                "cost_analysis",
                                "memory_analysis",
                                "runtime_executable"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "Lowered": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__",
                                "as_text",
                                "compile",
                                "compiler_ir",
                                "cost_analysis"
                            ],
                            "methods": [
                                "from_flat_info"
                            ],
                            "modules": {}
                        },
                        "Wrapped": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__call__",
                                "__init__",
                                "__new__",
                                "__subclasshook__",
                                "lower"
                            ],
                            "methods": [
                                "__class_getitem__",
                                "__init_subclass__"
                            ],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Compiled": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Lowered": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Wrapped": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [],
                "methods": [],
                "modules": {}
            },
            "jax.tree_util": {
                "classes": {
                    "AttributeKeyPathEntry": {
                        "AttributeKeyPathEntry": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__str__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace",
                                "pprint"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "DictKey": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "FlattenedIndexKey": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "GetAttrKey": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "GetitemKeyPathEntry": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__str__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace",
                                "pprint"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "Partial": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "PyTreeDef": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "SequenceKey": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__str__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "DictKey": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "FlattenedIndexKey": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "GetAttrKey": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "GetitemKeyPathEntry": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Partial": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PyTreeDef": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "SequenceKey": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "all_leaves",
                    "build_tree",
                    "keystr",
                    "register_keypaths",
                    "register_pytree_node",
                    "register_pytree_node_class",
                    "register_pytree_with_keys",
                    "register_pytree_with_keys_class",
                    "tree_all",
                    "tree_flatten",
                    "tree_flatten_with_path",
                    "tree_leaves",
                    "tree_map",
                    "tree_map_with_path",
                    "tree_reduce",
                    "tree_structure",
                    "tree_transpose",
                    "tree_unflatten",
                    "treedef_children",
                    "treedef_is_leaf",
                    "treedef_tuple"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.typing": {
                "classes": {},
                "functions": [],
                "methods": [],
                "modules": {}
            },
            "jax.util": {
                "classes": {
                    "HashableFunction": {
                        "HashableFunction": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__call__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "as_hashable_function",
                    "cache",
                    "safe_map",
                    "safe_zip",
                    "split_dict",
                    "split_list",
                    "split_merge",
                    "subvals",
                    "toposort",
                    "unzip2",
                    "wrap_name",
                    "wraps"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.version": {
                "classes": {},
                "functions": [
                    "_version_as_tuple"
                ],
                "methods": [],
                "modules": {}
            }
        }
    },
    "jax.abstract_arrays": {
        "classes": {
            "ShapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "raise_to_shaped"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.ad_checkpoint": {
        "classes": {},
        "functions": [
            "checkpoint",
            "checkpoint_name",
            "print_saved_residuals",
            "remat"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.api_util": {
        "classes": {},
        "functions": [
            "argnums_partial",
            "donation_vector",
            "flatten_axes",
            "rebase_donate_argnums",
            "safe_map",
            "shaped_abstractify"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.core": {
        "classes": {
            "AbstractToken": {
                "AbstractToken": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "AbstractValue": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "CallPrimitive": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "abstract_eval",
                        "bind",
                        "bind_with_trace",
                        "def_abstract_eval",
                        "def_custom_bind",
                        "def_effectful_abstract_eval",
                        "def_impl",
                        "get_bind_params",
                        "impl"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ClosedJaxpr": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__",
                        "_repr_pretty_",
                        "map_jaxpr",
                        "pretty_print",
                        "replace"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ConcretizationTypeError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DShapedArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__ne__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__",
                        "_add",
                        "_and",
                        "_array_module",
                        "_bool",
                        "_complex",
                        "_copy",
                        "_deepcopy",
                        "_div",
                        "_divmod",
                        "_eq",
                        "_float",
                        "_floordiv",
                        "_ge",
                        "_getitem",
                        "_gt",
                        "_hex",
                        "_int",
                        "_iter",
                        "_le",
                        "_lshift",
                        "_lt",
                        "_matmul",
                        "_mod",
                        "_mul",
                        "_ne",
                        "_nonzero",
                        "_oct",
                        "_or",
                        "_pow",
                        "_radd",
                        "_rand",
                        "_rdiv",
                        "_rdivmod",
                        "_rfloordiv",
                        "_rlshift",
                        "_rmatmul",
                        "_rmod",
                        "_rmul",
                        "_ror",
                        "_round",
                        "_rpow",
                        "_rrshift",
                        "_rshift",
                        "_rsub",
                        "_rtruediv",
                        "_rxor",
                        "_setitem",
                        "_sub",
                        "_truediv",
                        "_xor",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DropVar": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__ge__",
                        "__gt__",
                        "__init__",
                        "__init_subclass__",
                        "__le__",
                        "__lt__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Effect": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "EvalTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "full_raise",
                        "lift",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "InDBIdx": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__delattr__",
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "InconclusiveDimensionOperation": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JaxprEqn": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace",
                        "replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "JaxprPpContext": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JaxprPpSettings": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "JaxprTypeError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Literal": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "MainTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "with_cur_sublevel"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "MapPrimitive": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "abstract_eval",
                        "bind",
                        "bind_with_trace",
                        "def_abstract_eval",
                        "def_custom_bind",
                        "def_effectful_abstract_eval",
                        "def_impl",
                        "get_bind_params",
                        "impl",
                        "post_process",
                        "process"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "NameGatheringSubst": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "NamedShape": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__getitem__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__str__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "OutDBIdx": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__delattr__",
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Primitive": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "abstract_eval",
                        "bind",
                        "bind_with_trace",
                        "def_abstract_eval",
                        "def_custom_bind",
                        "def_effectful_abstract_eval",
                        "def_impl",
                        "get_bind_params",
                        "impl"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Sublevel": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__ge__",
                        "__gt__",
                        "__init__",
                        "__init_subclass__",
                        "__le__",
                        "__lt__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ThreadLocalState": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Token": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Trace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "full_raise",
                        "lift",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "TraceStack": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "copy",
                        "next_level",
                        "pop",
                        "push"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TraceState": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "copy"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Tracer": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__abs__",
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__ge__",
                        "__getattr__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__invert__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__long__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__neg__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pos__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_assert_live",
                        "_contents",
                        "_origin_msg",
                        "_pretty_print",
                        "addressable_data",
                        "delete",
                        "device",
                        "devices",
                        "get_referent",
                        "is_deleted",
                        "on_device_size_in_bytes",
                        "tobytes",
                        "tolist",
                        "unsafe_buffer_pointer"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TracerArrayConversionError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TracerIntegerConversionError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "UnexpectedTracerError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "UnshapedArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__ne__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_bool",
                        "_complex",
                        "_float",
                        "_hex",
                        "_int",
                        "_nonzero",
                        "_oct",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Var": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__ge__",
                        "__gt__",
                        "__init__",
                        "__init_subclass__",
                        "__le__",
                        "__lt__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": [],
                "partial": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "weakref": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                }
            },
            "AbstractValue": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "CallPrimitive": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ClosedJaxpr": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ConcreteArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ConcretizationTypeError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DShapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DropVar": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Effect": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "EvalTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "InDBIdx": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "InconclusiveDimensionOperation": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Jaxpr": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprEqn": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprPpContext": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprPpSettings": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprTypeError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Literal": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MainTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MapPrimitive": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "NameGatheringSubst": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "NamedShape": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "OutDBIdx": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Primitive": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Sublevel": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ThreadLocalState": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Token": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Trace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TraceStack": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TraceState": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Tracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerArrayConversionError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerIntegerConversionError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UnexpectedTracerError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UnshapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Var": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "partial": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ref": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "apply_todos",
            "as_hashable_function",
            "as_named_shape",
            "axis_frame",
            "call_bind_with_continuation",
            "call_impl",
            "canonicalize_dim",
            "canonicalize_shape",
            "check_eqn",
            "check_jaxpr",
            "check_type",
            "check_valid_jaxtype",
            "concrete_aval",
            "concrete_or_error",
            "concretization_function_error",
            "cur_sublevel",
            "dedup_referents",
            "diff_dim",
            "diff_shape",
            "dilate_dim",
            "dilate_shape",
            "dimension_as_value",
            "divide_shape_sizes",
            "do_subst_axis_names_jaxpr",
            "ensure_compile_time_eval",
            "escaped_tracer_error",
            "eval_context",
            "eval_jaxpr",
            "extend_axis_env",
            "extend_axis_env_nd",
            "find_top_trace",
            "full_lower",
            "gensym",
            "get_aval",
            "get_referent",
            "greater_equal_dim",
            "greater_equal_shape",
            "has_opaque_dtype",
            "is_constant_dim",
            "is_constant_shape",
            "is_dim",
            "is_empty_shape",
            "is_opaque_dtype",
            "is_special_dim_size",
            "jaxpr_uses_outfeed",
            "jaxprs_in_params",
            "join_effects",
            "join_named_shapes",
            "lattice_join",
            "leaked_tracer_error",
            "map",
            "map_bind",
            "map_bind_with_continuation",
            "mapped_aval",
            "maybe_find_leaked_tracers",
            "namedtuple",
            "new_base_main",
            "new_jaxpr_eqn",
            "new_main",
            "new_sublevel",
            "pp_aval",
            "pp_eqn",
            "pp_eqns",
            "pp_jaxpr",
            "pp_jaxpr_eqn_range",
            "pp_jaxpr_skeleton",
            "pp_jaxprs",
            "pp_kv_pair",
            "pp_kv_pairs",
            "pp_var",
            "pp_vars",
            "primal_dtype_to_tangent_dtype",
            "primitive_uses_outfeed",
            "raise_as_much_as_possible",
            "raise_to_shaped",
            "reset_trace_state",
            "safe_map",
            "safe_zip",
            "same_referent",
            "same_shape_sizes",
            "stash_axis_env",
            "str_eqn_compact",
            "stride_dim",
            "stride_shape",
            "subjaxprs",
            "subst_axis_names",
            "subst_axis_names_eqn",
            "subst_axis_names_jaxpr",
            "subst_axis_names_var",
            "substitute_vars_in_output_ty",
            "sum_dim",
            "sum_shapes",
            "symbolic_equal_dim",
            "symbolic_equal_one_of_dim",
            "symbolic_equal_shape",
            "total_ordering",
            "trace_state_clean",
            "traverse_jaxpr_params",
            "tuple_delete",
            "tuple_insert",
            "typecheck",
            "typecompat",
            "typematch",
            "unmapped_aval",
            "used_axis_names",
            "valid_jaxtype",
            "zip"
        ],
        "methods": [
            "call"
        ],
        "modules": {}
    },
    "jax.custom_batching": {
        "classes": {
            "custom_vmap": {
                "classes": {},
                "custom_vmap": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__getattr__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "def_vmap"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "sequential_vmap"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.custom_derivatives": {
        "classes": {
            "SymbolicZero": {
                "SymbolicZero": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__getattr__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "custom_jvp": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__getattr__",
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "defjvp",
                        "defjvps"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "custom_vjp": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__getattr__",
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "defvjp"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "functions": [],
                "methods": []
            },
            "custom_jvp": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "custom_vjp": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "_initial_style_jaxpr",
            "_sum_tangents",
            "_zeros_like_pytree",
            "closure_convert",
            "custom_gradient",
            "linear_call"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.custom_transpose": {
        "classes": {
            "custom_transpose": {
                "classes": {},
                "custom_transpose": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__getattr__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "def_transpose"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "functions": [],
                "methods": []
            }
        },
        "functions": [],
        "methods": [],
        "modules": {}
    },
    "jax.debug": {
        "classes": {
            "DebugEffect": {
                "DebugEffect": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__str__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "breakpoint",
            "callback",
            "inspect_array_sharding",
            "print",
            "visualize_array_sharding",
            "visualize_sharding"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.distributed": {
        "classes": {},
        "functions": [
            "initialize",
            "shutdown"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.dtypes": {
        "classes": {
            "bfloat16": {
                "bfloat16": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "finfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__",
                        "_bfloat16_finfo",
                        "_float8_e4m3fn_finfo",
                        "_float8_e5m2_finfo",
                        "_init"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "functions": [],
                "iinfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "methods": []
            },
            "finfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "iinfo": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "canonicalize_dtype",
            "issubdtype",
            "result_type",
            "scalar_type_of"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.errors": {
        "classes": {
            "ConcretizationTypeError": {
                "JAXIndexError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JAXTypeError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "NonConcreteBooleanIndexError": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JAXIndexError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JAXTypeError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "NonConcreteBooleanIndexError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerArrayConversionError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerIntegerConversionError": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UnexpectedTracerError": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [],
        "methods": [],
        "modules": {}
    },
    "jax.image": {
        "classes": {
            "ResizeMethod": {
                "ResizeMethod": {
                    "classes": {
                        "__class__": {
                            "EnumMeta": {
                                "classes": {
                                    "__base__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    },
                                    "__class__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    }
                                },
                                "functions": [
                                    "__bool__",
                                    "__call__",
                                    "__contains__",
                                    "__delattr__",
                                    "__dir__",
                                    "__getattr__",
                                    "__getitem__",
                                    "__init_subclass__",
                                    "__iter__",
                                    "__len__",
                                    "__new__",
                                    "__repr__",
                                    "__reversed__",
                                    "__setattr__",
                                    "__subclasshook__",
                                    "_check_for_existing_members",
                                    "_convert_",
                                    "_create_",
                                    "_find_new_",
                                    "_get_mixins_"
                                ],
                                "methods": [
                                    "__prepare__"
                                ],
                                "modules": {}
                            },
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "resize",
            "scale_and_translate"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.interpreters": {
        "classes": {},
        "functions": [],
        "methods": [],
        "modules": {
            "jax.interpreters.batching": {
                "classes": {
                    "BatchTrace": {
                        "BatchTrace": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "full_raise",
                                "get_axis_primitive_batcher",
                                "get_frame",
                                "get_primitive_batcher",
                                "lift",
                                "post_process_call",
                                "post_process_custom_jvp_call",
                                "post_process_custom_vjp_call",
                                "post_process_custom_vjp_call_fwd",
                                "post_process_map",
                                "process_axis_index",
                                "process_call",
                                "process_custom_jvp_call",
                                "process_custom_transpose",
                                "process_custom_vjp_call",
                                "process_map",
                                "process_primitive",
                                "pure",
                                "sublift"
                            ],
                            "methods": [
                                "__class_getitem__",
                                "__init_subclass__"
                            ],
                            "modules": {}
                        },
                        "BatchTracer": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__abs__",
                                "__add__",
                                "__and__",
                                "__array__",
                                "__array_module__",
                                "__bool__",
                                "__complex__",
                                "__copy__",
                                "__deepcopy__",
                                "__div__",
                                "__divmod__",
                                "__dlpack__",
                                "__eq__",
                                "__float__",
                                "__floordiv__",
                                "__ge__",
                                "__getattr__",
                                "__getitem__",
                                "__gt__",
                                "__hex__",
                                "__index__",
                                "__init__",
                                "__init_subclass__",
                                "__int__",
                                "__invert__",
                                "__iter__",
                                "__le__",
                                "__len__",
                                "__long__",
                                "__lshift__",
                                "__lt__",
                                "__matmul__",
                                "__mod__",
                                "__mul__",
                                "__ne__",
                                "__neg__",
                                "__new__",
                                "__nonzero__",
                                "__oct__",
                                "__or__",
                                "__pos__",
                                "__pow__",
                                "__radd__",
                                "__rand__",
                                "__rdiv__",
                                "__rdivmod__",
                                "__reduce__",
                                "__repr__",
                                "__reversed__",
                                "__rfloordiv__",
                                "__rlshift__",
                                "__rmatmul__",
                                "__rmod__",
                                "__rmul__",
                                "__ror__",
                                "__round__",
                                "__rpow__",
                                "__rrshift__",
                                "__rshift__",
                                "__rsub__",
                                "__rtruediv__",
                                "__rxor__",
                                "__setitem__",
                                "__sub__",
                                "__subclasshook__",
                                "__truediv__",
                                "__xor__",
                                "_assert_live",
                                "_contents",
                                "_origin_msg",
                                "_pretty_print",
                                "addressable_data",
                                "delete",
                                "device",
                                "devices",
                                "full_lower",
                                "get_referent",
                                "is_deleted",
                                "on_device_size_in_bytes",
                                "tobytes",
                                "tolist",
                                "unsafe_buffer_pointer"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "ConcatAxis": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "IndexedAxisSize": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__",
                                "replace"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "NoneType": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "Pile": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "PileAxis": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "PileTy": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__",
                                "replace"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "ZeroIfMapped": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init_subclass__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "BatchTracer": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ConcatAxis": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "IndexedAxisSize": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "NotMapped": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Pile": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PileAxis": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "PileTy": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "Zero": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ZeroIfMapped": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "add_batched",
                    "batch",
                    "batch_custom_vjp_bwd",
                    "batch_jaxpr",
                    "batch_jaxpr2",
                    "batch_jaxpr_axes",
                    "bdim_at_front",
                    "broadcast",
                    "broadcast_batcher",
                    "defbroadcasting",
                    "defreducer",
                    "defvectorized",
                    "from_elt",
                    "is_vmappable",
                    "make_iota",
                    "matchaxis",
                    "moveaxis",
                    "reassemble_concat_axes",
                    "reducer_batcher",
                    "register_vmappable",
                    "segment_sum",
                    "to_elt",
                    "unpack_concat_axes",
                    "unregister_vmappable",
                    "vectorized_batcher",
                    "vtile",
                    "zeros_like_batched"
                ],
                "methods": [],
                "modules": {}
            },
            "jax.interpreters.mlir": {
                "classes": {
                    "ConstantHandler": {
                        "ConstantHandler": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__call__",
                                "__init__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [
                                "__class_getitem__",
                                "__init_subclass__"
                            ],
                            "modules": {}
                        },
                        "DimExprEvaluator": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__add__",
                                "__divmod__",
                                "__init__",
                                "__init_subclass__",
                                "__mul__",
                                "__new__",
                                "__radd__",
                                "__rdivmod__",
                                "__rmul__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "LoweringResult": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__class_getitem__",
                                "__getnewargs__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "_asdict",
                                "_replace"
                            ],
                            "methods": [
                                "_make"
                            ],
                            "modules": {}
                        },
                        "LoweringRule": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__call__",
                                "__init__",
                                "__new__",
                                "__subclasshook__"
                            ],
                            "methods": [
                                "__class_getitem__",
                                "__init_subclass__"
                            ],
                            "modules": {}
                        },
                        "LoweringRuleContext": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "replace",
                                "set_tokens_out"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "ModuleContext": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__eq__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__subclasshook__",
                                "add_host_callback",
                                "add_keepalive",
                                "new_channel",
                                "replace"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "ReplicaAxisContext": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "SPMDAxisContext": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__",
                                "extend_manual"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "ShardingContext": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__delattr__",
                                "__eq__",
                                "__hash__",
                                "__init__",
                                "__init_subclass__",
                                "__new__",
                                "__repr__",
                                "__setattr__",
                                "__subclasshook__"
                            ],
                            "methods": [],
                            "modules": {}
                        },
                        "TokenSet": {
                            "classes": {
                                "__class__": {
                                    "classes": {},
                                    "functions": [],
                                    "methods": []
                                }
                            },
                            "functions": [
                                "__init__",
                                "__init_subclass__",
                                "__len__",
                                "__new__",
                                "__subclasshook__",
                                "effects",
                                "get",
                                "items",
                                "subset",
                                "update_tokens"
                            ],
                            "methods": [
                                "create"
                            ],
                            "modules": {}
                        },
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "DimExprEvaluator": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "LoweringResult": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "LoweringRule": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "LoweringRuleContext": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ModuleContext": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ReplicaAxisContext": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "SPMDAxisContext": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "ShardingContext": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    },
                    "TokenSet": {
                        "classes": {},
                        "functions": [],
                        "methods": []
                    }
                },
                "functions": [
                    "_call_lowering",
                    "aval_to_ir_type",
                    "aval_to_ir_types",
                    "core_call_lowering",
                    "dense_bool_elements",
                    "dense_int_elements",
                    "dtype_to_ir_type",
                    "emit_python_callback",
                    "flatten_lowering_ir_args",
                    "i32_attr",
                    "i64_attr",
                    "ir_constant",
                    "ir_constants",
                    "jaxpr_subcomp",
                    "lower_fun",
                    "lower_jaxpr_to_fun",
                    "lower_jaxpr_to_module",
                    "make_ir_context",
                    "merge_mlir_modules",
                    "module_to_bytecode",
                    "module_to_string",
                    "register_constant_handler",
                    "register_lowering",
                    "shape_tensor",
                    "token_type",
                    "xla_computation_to_mlir_module"
                ],
                "methods": [],
                "modules": {}
            }
        }
    },
    "jax.interpreters.ad": {
        "classes": {
            "CustomJVPException": {
                "CustomJVPException": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": [],
                            "type": {
                                "classes": {
                                    "__base__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": [],
                                        "object": {
                                            "classes": {
                                                "__class__": {
                                                    "classes": {},
                                                    "functions": [],
                                                    "methods": []
                                                }
                                            },
                                            "functions": [
                                                "__init_subclass__",
                                                "__new__",
                                                "__subclasshook__"
                                            ],
                                            "methods": [],
                                            "modules": {}
                                        }
                                    },
                                    "__class__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    }
                                },
                                "functions": [
                                    "__init_subclass__",
                                    "__new__",
                                    "__prepare__",
                                    "__subclasshook__"
                                ],
                                "methods": [],
                                "modules": {}
                            }
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "CustomVJPException": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JVPTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "full_raise",
                        "join",
                        "lift",
                        "post_process_call",
                        "post_process_custom_jvp_call",
                        "post_process_custom_vjp_call",
                        "post_process_map",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "JVPTracer": {
                    "classes": {
                        "__class__": {
                            "ABCMeta": {
                                "classes": {
                                    "__base__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    },
                                    "__class__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    }
                                },
                                "functions": [
                                    "__init_subclass__",
                                    "__instancecheck__",
                                    "__new__",
                                    "__prepare__",
                                    "__subclasscheck__",
                                    "__subclasshook__",
                                    "_abc_caches_clear",
                                    "_abc_registry_clear",
                                    "_dump_registry",
                                    "register"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__abs__",
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__ge__",
                        "__getattr__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__invert__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__long__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__neg__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pos__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_assert_live",
                        "_contents",
                        "_origin_msg",
                        "_pretty_print",
                        "addressable_data",
                        "delete",
                        "device",
                        "devices",
                        "full_lower",
                        "get_referent",
                        "is_deleted",
                        "on_device_size_in_bytes",
                        "tobytes",
                        "tolist",
                        "unsafe_buffer_pointer"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "UndefinedPrimal": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Zero": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "from_value"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            },
            "CustomVJPException": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JVPTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JVPTracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UndefinedPrimal": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Zero": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "add_jaxvals",
            "add_tangents",
            "backward_pass",
            "bilinear_transpose",
            "call_transpose",
            "closed_backward_pass",
            "defbilinear",
            "defjvp",
            "defjvp2",
            "defjvp_zero",
            "deflinear",
            "deflinear2",
            "get_primitive_transpose",
            "instantiate_zeros",
            "instantiate_zeros_aval",
            "is_undefined_primal",
            "jvp",
            "jvp_jaxpr",
            "linear_jvp",
            "linear_transpose",
            "linear_transpose2",
            "linearize",
            "map_transpose",
            "rearrange_binders",
            "recast_to_float0",
            "replace_float0s",
            "standard_jvp",
            "standard_jvp2",
            "unpair_pval",
            "vjp",
            "zero_jvp",
            "zeros_like_aval",
            "zeros_like_jaxval"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.interpreters.partial_eval": {
        "classes": {
            "BoundedAxisSize": {
                "BoundedAxisSize": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "ConstVar": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "DebugInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "DimensionHandlerTracer": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "as_value",
                        "diff",
                        "dilate",
                        "divide_shape_sizes",
                        "greater_equal",
                        "is_constant",
                        "stride",
                        "sum",
                        "symbolic_equal"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DynamicJaxprTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_lift_tracers_in_aval",
                        "_new_const",
                        "default_process_primitive",
                        "full_raise",
                        "getvar",
                        "instantiate_const",
                        "lift",
                        "makevar",
                        "new_arg",
                        "new_const",
                        "post_process_call",
                        "post_process_custom_jvp_call",
                        "post_process_custom_vjp_call",
                        "post_process_map",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "DynamicJaxprTracer": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__abs__",
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__ge__",
                        "__getattr__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__invert__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__long__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__neg__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pos__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_assert_live",
                        "_contents",
                        "_origin_msg",
                        "_pretty_print",
                        "addressable_data",
                        "delete",
                        "device",
                        "devices",
                        "full_lower",
                        "get_referent",
                        "is_deleted",
                        "on_device_size_in_bytes",
                        "tobytes",
                        "tolist",
                        "unsafe_buffer_pointer"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "FreeVar": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "Jaxpr": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__str__",
                        "__subclasshook__",
                        "_repr_pretty_",
                        "pretty_print",
                        "replace"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JaxprEqnRecipe": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "JaxprStackFrame": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "add_eqn",
                        "find_progenitors",
                        "newvar",
                        "to_jaxpr",
                        "to_jaxpr2"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "JaxprTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_current_truncated_name_stack",
                        "default_process_primitive",
                        "full_raise",
                        "instantiate_const",
                        "instantiate_const_abstracted",
                        "lift",
                        "new_arg",
                        "new_const",
                        "new_instantiated_const",
                        "new_instantiated_literal",
                        "post_process_call",
                        "post_process_custom_jvp_call",
                        "post_process_custom_vjp_call",
                        "post_process_map",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "JaxprTracer": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__abs__",
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__ge__",
                        "__getattr__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__invert__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__long__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__neg__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pos__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_assert_live",
                        "_contents",
                        "_origin_msg",
                        "_pretty_print",
                        "addressable_data",
                        "delete",
                        "device",
                        "devices",
                        "full_lower",
                        "get_referent",
                        "is_deleted",
                        "is_known",
                        "on_device_size_in_bytes",
                        "tobytes",
                        "tolist",
                        "unsafe_buffer_pointer"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "LambdaBinding": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "PartialVal": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "get_aval",
                        "get_known",
                        "is_known"
                    ],
                    "methods": [
                        "known",
                        "unknown"
                    ],
                    "modules": {}
                },
                "TracerAsName": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "int": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "from_bytes"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "methods": []
            },
            "ConstVar": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DebugInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DimensionHandlerTracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DynamicJaxprTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DynamicJaxprTracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "FreeVar": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Jaxpr": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprEqnRecipe": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprStackFrame": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "JaxprTracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "LambdaBinding": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "PartialVal": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerAsName": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TracerId": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "abstract_eval_fun",
            "arg_info_all",
            "call_padding_rule",
            "call_partial_eval_custom_rule",
            "closed_call_partial_eval_custom_rule",
            "convert_envvars_to_constvars",
            "dce_jaxpr",
            "dce_jaxpr_call_rule",
            "dce_jaxpr_closed_call_rule",
            "dce_jaxpr_consts",
            "debug_info",
            "debug_info_final",
            "def_trivial_padding",
            "extend_jaxpr_stack",
            "infer_lambda_input_type",
            "instantiate_const_at",
            "make_jaxpr_effects",
            "move_binders_to_back",
            "move_binders_to_front",
            "new_eqn_recipe",
            "pad_jaxpr",
            "partial_eval_jaxpr_custom",
            "partial_eval_jaxpr_custom_rule_not_implemented",
            "partial_eval_jaxpr_nounits",
            "partition_pvals",
            "recipe_to_eqn",
            "result_info",
            "sig_info",
            "trace_to_jaxpr",
            "trace_to_jaxpr_dynamic",
            "trace_to_jaxpr_dynamic2",
            "trace_to_jaxpr_final",
            "trace_to_jaxpr_final2",
            "trace_to_jaxpr_nounits",
            "trace_to_subjaxpr_dynamic",
            "trace_to_subjaxpr_dynamic2",
            "tracers_to_jaxpr",
            "trivial_ctx"
        ],
        "methods": [],
        "modules": {}
    },
    "jax.interpreters.pxla": {
        "classes": {
            "ArrayMapping": {
                "Chunked": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": [],
                            "pybind11_type": {
                                "classes": {
                                    "__base__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    },
                                    "__class__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    }
                                },
                                "functions": [
                                    "__init_subclass__",
                                    "__new__",
                                    "__prepare__",
                                    "__subclasshook__"
                                ],
                                "methods": [],
                                "modules": {}
                            }
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DynamicAxisEnv": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__contains__",
                        "__getitem__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DynamicAxisEnvFrame": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "EmapInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "ExecuteReplicated": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__",
                        "_add_tokens_to_inputs",
                        "_call_with_tokens",
                        "_handle_token_bufs"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "InputsHandler": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__str__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "MapTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "full_raise",
                        "lift",
                        "process_axis_index",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "MapTracer": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__abs__",
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__ge__",
                        "__getattr__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__invert__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__long__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__neg__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pos__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__str__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_assert_live",
                        "_contents",
                        "_origin_msg",
                        "_pretty_print",
                        "addressable_data",
                        "delete",
                        "device",
                        "devices",
                        "full_lower",
                        "get_referent",
                        "is_deleted",
                        "on_device_size_in_bytes",
                        "tobytes",
                        "tolist",
                        "unsafe_buffer_pointer"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Mesh": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__enter__",
                        "__eq__",
                        "__exit__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__",
                        "_local_mesh",
                        "_recreate_cm"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "MeshComputation": {
                    "classes": {
                        "__class__": {
                            "_ProtocolMeta": {
                                "classes": {
                                    "__base__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    },
                                    "__class__": {
                                        "classes": {},
                                        "functions": [],
                                        "methods": []
                                    }
                                },
                                "functions": [
                                    "__init_subclass__",
                                    "__instancecheck__",
                                    "__new__",
                                    "__prepare__",
                                    "__subclasscheck__",
                                    "__subclasshook__",
                                    "_abc_caches_clear",
                                    "_abc_registry_clear",
                                    "_dump_registry",
                                    "register"
                                ],
                                "methods": [],
                                "modules": {}
                            },
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "as_text",
                        "compile",
                        "compiler_ir",
                        "cost_analysis",
                        "hlo",
                        "mhlo",
                        "stablehlo"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "MeshExecutable": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "as_text",
                        "call",
                        "cost_analysis",
                        "create_cpp_call",
                        "from_trivial_jaxpr",
                        "input_shardings",
                        "memory_analysis",
                        "output_shardings",
                        "runtime_executable",
                        "xla_extension_executable"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "NoSharding": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "OrderedDict": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        },
                        "__origin__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init_subclass__",
                        "__mro_entries__",
                        "__new__",
                        "__reduce__",
                        "__reduce_ex__",
                        "__subclasshook__",
                        "fromkeys"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ParallelCallableInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__delattr__",
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "PartitionInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "PartitionSpec": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__reduce__",
                        "__repr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "PmapComputation": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "as_text",
                        "compile",
                        "compiler_ir",
                        "cost_analysis",
                        "hlo",
                        "mhlo",
                        "stablehlo"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "PmapExecutable": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__subclasshook__",
                        "as_text",
                        "call",
                        "cost_analysis",
                        "create_cpp_call",
                        "input_shardings",
                        "memory_analysis",
                        "output_shardings",
                        "runtime_executable",
                        "xla_extension_executable"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "ReplicaInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "Replicated": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ResultsHandler": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "SPMDBatchTrace": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "full_raise",
                        "get_axis_primitive_batcher",
                        "get_frame",
                        "get_primitive_batcher",
                        "lift",
                        "post_process_call",
                        "post_process_custom_jvp_call",
                        "post_process_custom_vjp_call",
                        "post_process_custom_vjp_call_fwd",
                        "post_process_map",
                        "process_axis_index",
                        "process_call",
                        "process_custom_jvp_call",
                        "process_custom_transpose",
                        "process_custom_vjp_call",
                        "process_map",
                        "process_primitive",
                        "pure",
                        "sublift"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "ShardInfo": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "ShardedAxis": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ShardedDeviceArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ShardingSpec": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "indices",
                        "sharding_proto"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TileManual": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__delattr__",
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TileVectorize": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__delattr__",
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__setattr__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "UnloadedMeshExecutable": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "build_unsafe_call",
                        "from_hlo",
                        "load"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "UnloadedPmapExecutable": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "build_execute_fun",
                        "from_hlo",
                        "load"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Unstacked": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "WeakRefList": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Chunked": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DynamicAxisEnv": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "DynamicAxisEnvFrame": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "EmapInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ExecuteReplicated": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "InputsHandler": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MapTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MapTracer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MeshComputation": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "MeshExecutable": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "NoSharding": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "OrderedDictType": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ParallelCallableInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "PartitionInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "PmapComputation": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "PmapExecutable": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ReplicaInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Replicated": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ResultsHandler": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "SPMDBatchTrace": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShardInfo": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShardedAxis": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShardingSpec": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TileManual": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TileVectorize": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UnloadedMeshExecutable": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "UnloadedPmapExecutable": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Unstacked": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "WeakRefList": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_Mesh": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_PartitionSpec": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_ShardedDeviceArray": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "__getattr__",
            "_create_pmap_sharding_spec",
            "_deprecated_device_put",
            "_deprecated_make_sharded_device_array",
            "_get_and_check_device_assignment",
            "_get_sharding_specs",
            "_is_unspecified",
            "_one_replica_buffer_indices",
            "_pmap_sharding_spec",
            "are_op_shardings_equal",
            "array_mapping_to_axis_resources",
            "find_partitions",
            "find_replicas",
            "get_global_aval",
            "get_local_aval",
            "get_num_partitions",
            "global_aval_to_result_handler",
            "global_avals_to_results_handler",
            "is_op_sharding_replicated",
            "local_aval_to_result_handler",
            "local_avals_to_results_handler",
            "lower_mesh_computation",
            "lower_parallel_callable",
            "lower_sharding_computation",
            "maybe_extend_axis_env",
            "mesh_sharding_specs",
            "new_mesh_sharding_specs",
            "op_sharding_to_indices",
            "parallel_callable",
            "partitioned_sharding_spec",
            "reconcile_num_partitions",
            "replicate",
            "resource_typecheck",
            "shard_arg",
            "shard_args",
            "shard_aval",
            "sharding_spec_sharding_proto",
            "spec_to_indices",
            "stage_parallel_callable",
            "tile_aval_nd",
            "untile_aval_nd",
            "vtile_by_mesh",
            "wrap_name",
            "xla_pmap_impl",
            "xla_pmap_impl_lazy"
        ],
        "methods": [
            "xla_pmap"
        ],
        "modules": {}
    },
    "jax.interpreters.xla": {
        "classes": {
            "AxisEnv": {
                "AxisEnv": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__class_getitem__",
                        "__getnewargs__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_asdict",
                        "_replace"
                    ],
                    "methods": [
                        "_make"
                    ],
                    "modules": {}
                },
                "Client": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ConcreteArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__ne__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_add",
                        "_and",
                        "_array_module",
                        "_bool",
                        "_complex",
                        "_copy",
                        "_deepcopy",
                        "_div",
                        "_divmod",
                        "_eq",
                        "_float",
                        "_floordiv",
                        "_ge",
                        "_getitem",
                        "_gt",
                        "_hex",
                        "_int",
                        "_iter",
                        "_le",
                        "_len",
                        "_lshift",
                        "_lt",
                        "_matmul",
                        "_mod",
                        "_mul",
                        "_ne",
                        "_nonzero",
                        "_oct",
                        "_or",
                        "_pow",
                        "_radd",
                        "_rand",
                        "_rdiv",
                        "_rdivmod",
                        "_rfloordiv",
                        "_rlshift",
                        "_rmatmul",
                        "_rmod",
                        "_rmul",
                        "_ror",
                        "_round",
                        "_rpow",
                        "_rrshift",
                        "_rshift",
                        "_rsub",
                        "_rtruediv",
                        "_rxor",
                        "_setitem",
                        "_sub",
                        "_truediv",
                        "_xor",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "Device": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DeviceArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__format__",
                        "__ge__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init_subclass__",
                        "__int__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__str__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_chunk_iter",
                        "_multi_slice",
                        "_split",
                        "all",
                        "any",
                        "argmax",
                        "argmin",
                        "astype",
                        "block_until_ready",
                        "broadcast",
                        "broadcast_in_dim",
                        "choose",
                        "clip",
                        "clone",
                        "compress",
                        "copy",
                        "copy_to_device",
                        "copy_to_host_async",
                        "copy_to_remote_device",
                        "cumprod",
                        "cumsum",
                        "delete",
                        "device",
                        "is_deleted",
                        "is_known_ready",
                        "is_ready",
                        "item",
                        "max",
                        "mean",
                        "min",
                        "nonzero",
                        "on_device_size_in_bytes",
                        "platform",
                        "prod",
                        "ptp",
                        "repeat",
                        "reshape",
                        "split",
                        "squeeze",
                        "std",
                        "sum",
                        "take",
                        "to_py",
                        "tobytes",
                        "tolist",
                        "transpose",
                        "unsafe_buffer_pointer",
                        "var",
                        "view",
                        "xla_dynamic_shape",
                        "xla_shape"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "DeviceArrayBase": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__add__",
                        "__and__",
                        "__array__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__format__",
                        "__ge__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init_subclass__",
                        "__int__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__str__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_chunk_iter",
                        "_multi_slice",
                        "_split",
                        "all",
                        "any",
                        "argmax",
                        "argmin",
                        "astype",
                        "broadcast",
                        "broadcast_in_dim",
                        "choose",
                        "clip",
                        "compress",
                        "copy",
                        "cumprod",
                        "cumsum",
                        "item",
                        "max",
                        "mean",
                        "min",
                        "nonzero",
                        "prod",
                        "ptp",
                        "repeat",
                        "reshape",
                        "split",
                        "squeeze",
                        "std",
                        "sum",
                        "take",
                        "to_py",
                        "tobytes",
                        "tolist",
                        "transpose",
                        "var",
                        "view"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "ShapedArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__hash__",
                        "__init__",
                        "__init_subclass__",
                        "__ne__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "_add",
                        "_and",
                        "_array_module",
                        "_bool",
                        "_complex",
                        "_copy",
                        "_deepcopy",
                        "_div",
                        "_divmod",
                        "_eq",
                        "_float",
                        "_floordiv",
                        "_ge",
                        "_getitem",
                        "_gt",
                        "_hex",
                        "_int",
                        "_iter",
                        "_le",
                        "_len",
                        "_lshift",
                        "_lt",
                        "_matmul",
                        "_mod",
                        "_mul",
                        "_ne",
                        "_nonzero",
                        "_oct",
                        "_or",
                        "_pow",
                        "_radd",
                        "_rand",
                        "_rdiv",
                        "_rdivmod",
                        "_rfloordiv",
                        "_rlshift",
                        "_rmatmul",
                        "_rmod",
                        "_rmul",
                        "_ror",
                        "_round",
                        "_rpow",
                        "_rrshift",
                        "_rshift",
                        "_rsub",
                        "_rtruediv",
                        "_rxor",
                        "_setitem",
                        "_sub",
                        "_truediv",
                        "_xor",
                        "at_least_vspace",
                        "join",
                        "str_short",
                        "strip_named_shape",
                        "strip_weak_type",
                        "update"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TranslationContext": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__eq__",
                        "__init__",
                        "__init_subclass__",
                        "__new__",
                        "__repr__",
                        "__subclasshook__",
                        "replace"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "TranslationRule": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__call__",
                        "__init__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [
                        "__class_getitem__",
                        "__init_subclass__"
                    ],
                    "modules": {}
                },
                "XlaOp": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__init_subclass__",
                        "__new__",
                        "__subclasshook__"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "_DeviceArray": {
                    "classes": {
                        "__class__": {
                            "classes": {},
                            "functions": [],
                            "methods": []
                        }
                    },
                    "functions": [
                        "__add__",
                        "__and__",
                        "__array__",
                        "__array_module__",
                        "__bool__",
                        "__complex__",
                        "__copy__",
                        "__deepcopy__",
                        "__div__",
                        "__divmod__",
                        "__dlpack__",
                        "__eq__",
                        "__float__",
                        "__floordiv__",
                        "__format__",
                        "__ge__",
                        "__getitem__",
                        "__gt__",
                        "__hex__",
                        "__index__",
                        "__init__",
                        "__init_subclass__",
                        "__int__",
                        "__iter__",
                        "__le__",
                        "__len__",
                        "__lshift__",
                        "__lt__",
                        "__matmul__",
                        "__mod__",
                        "__mul__",
                        "__ne__",
                        "__new__",
                        "__nonzero__",
                        "__oct__",
                        "__or__",
                        "__pow__",
                        "__radd__",
                        "__rand__",
                        "__rdiv__",
                        "__rdivmod__",
                        "__reduce__",
                        "__repr__",
                        "__reversed__",
                        "__rfloordiv__",
                        "__rlshift__",
                        "__rmatmul__",
                        "__rmod__",
                        "__rmul__",
                        "__ror__",
                        "__round__",
                        "__rpow__",
                        "__rrshift__",
                        "__rshift__",
                        "__rsub__",
                        "__rtruediv__",
                        "__rxor__",
                        "__setitem__",
                        "__str__",
                        "__sub__",
                        "__subclasshook__",
                        "__truediv__",
                        "__xor__",
                        "_check_if_deleted",
                        "_chunk_iter",
                        "_multi_slice",
                        "_split",
                        "all",
                        "any",
                        "argmax",
                        "argmin",
                        "astype",
                        "block_until_ready",
                        "broadcast",
                        "broadcast_in_dim",
                        "choose",
                        "clip",
                        "compress",
                        "copy",
                        "copy_to_host_async",
                        "cumprod",
                        "cumsum",
                        "delete",
                        "device",
                        "item",
                        "max",
                        "mean",
                        "min",
                        "nonzero",
                        "prod",
                        "ptp",
                        "repeat",
                        "reshape",
                        "split",
                        "squeeze",
                        "std",
                        "sum",
                        "take",
                        "to_py",
                        "tobytes",
                        "tolist",
                        "transpose",
                        "unsafe_buffer_pointer",
                        "var",
                        "view"
                    ],
                    "methods": [],
                    "modules": {}
                },
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Backend": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "Buffer": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ConcreteArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "ShapedArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TranslationContext": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "TranslationRule": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "XlaOp": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_CppDeviceArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_DeviceArray": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_Device": {
                "classes": {},
                "functions": [],
                "methods": []
            },
            "_deprecated_DeviceArray": {
                "classes": {},
                "functions": [],
                "methods": []
            }
        },
        "functions": [
            "__getattr__",
            "_deprecated_device_put",
            "abstractify",
            "apply_primitive",
            "axis_groups",
            "backend_compile",
            "canonicalize_dtype",
            "check_backend_matches",
            "make_device_array",
            "parameter",
            "register_collective_primitive",
            "register_initial_style_primitive",
            "register_translation",
            "sharding_to_proto",
            "xla_destructure"
        ],
        "methods": [],
        "modules": {}
    }
}